﻿#include "RemapUserObject.h"
#include "Application.h"
#include "Mesh.h"
#include "SGHSystem.h"
#include "Timer.h"
#include "Transient.h"

#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/ArrayCopy.h>
#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/cont/ArrayHandleGroupVecVariable.h>
#include <vtkm/io/VTKDataSetWriter.h>
#include <vtkm/worklet/ScatterUniform.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/cont/CellLocatorBoundingIntervalHierarchy.h>
#include <vtkm/cont/CellLocatorGeneral.h>
#include <vtkm/cont/CellLocatorTwoLevel.h>
#include <vtkm/exec/CellInterpolate.h>

class FindCellWorklet : public vtkm::worklet::WorkletMapField
{
public:
  using ControlSignature = void(FieldIn points,
                                ExecObject locator,
                                WholeCellSetIn<Cell, Point> cell_set,
                                WholeArrayIn field_in,
                                FieldOut field_out);
  using ExecutionSignature = void(_1, _2, _3, _4, _5);

  template<typename LocatorType, typename CellSet, typename FieldPortal, typename OutType>
  VTKM_EXEC void operator()(const vtkm::Vec3f& point,
                            const LocatorType& locator,
                            const CellSet& cell_set,
                            const FieldPortal& field_in,
                            OutType& out) const
  {
    Id cell_id;
    Vec3f pcoords;
    vtkm::ErrorCode status = locator.FindCell(0.99999 * point, cell_id, pcoords);

    if (cell_id >= 0)
    {
      auto cellShape = cell_set.GetCellShape(cell_id);
      auto indices = cell_set.GetIndices(cell_id);
      auto fieldValues = vtkm::make_VecFromPortalPermute(&indices, &field_in);
      vtkm::exec::CellInterpolate(fieldValues, pcoords, cellShape, out);
    }
    else
    {
      //console::Error(point);
      this->RaiseError(" Given point outside of the cell set .");
    }
  }
};

struct InterpolateVelocity : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn points,
                                ExecObject locator,
                                WholeCellSetIn<Cell, Point> cell_set,
                                WholeArrayIn field_in,
                                FieldOut field_out);
  using ExecutionSignature = void(_1, _2, _3, _4, _5);

  template<typename LocatorType, typename CellSet, typename FieldPortal, typename OutType>
  VTKM_EXEC void operator()(const vtkm::Vec3f& point,
                            const LocatorType& locator,
                            const CellSet& cell_set,
                            const FieldPortal& field_in,
                            OutType& out) const
  {
    Id cell_id;
    Vec3f pcoords;
    //vtkm::ErrorCode status = locator.FindCell(0.99999 * point, cell_id, pcoords);
    for (size_t i = 0; i < 5; i++)
    {
      Vec3f p = 0.99999 * point + offsets[i];
      vtkm::ErrorCode status = locator.FindCell(p, cell_id, pcoords);
      if (cell_id >= 0)
      {
        auto cellShape = cell_set.GetCellShape(cell_id);
        auto indices = cell_set.GetIndices(cell_id);
        auto fieldValues = vtkm::make_VecFromPortalPermute(&indices, &field_in);
        vtkm::exec::CellInterpolate(fieldValues, pcoords, cellShape, out);
        //console::Success(point);

        break;
      }
      //console::Error(point);
      //this->RaiseError(" Given point outside of the cell set .");
    }
    //if (cell_id >= 0)
    //{
    //  auto cellShape = cell_set.GetCellShape(cell_id);
    //  auto indices = cell_set.GetIndices(cell_id);
    //  auto fieldValues = vtkm::make_VecFromPortalPermute(&indices, &field_in);
    //  vtkm::exec::CellInterpolate(fieldValues, pcoords, cellShape, out);
    //}
    //else
    //{
    //  Real x = point[0];
    //  Real y = point[1];
    //  Vec3f p = point;
    //  if (x < 0.5 && y < 0.5) // bottom
    //  {
    //    p += Vec3f{ 0, 1, 0 };
    //  }
    //  else if (x > 0.5 && y > 0.5) // top
    //  {
    //    p += Vec3f{ 0, -1, 0 };
    //  }
    //  else if (x < 0.5 && y > 0.5) // left
    //  {
    //    p += Vec3f{ 1, 0, 0 };
    //  }
    //  else //if (x > 0.5 && y < 0.5) // right
    //  {
    //    p += Vec3f{ -1, 0, 0 };
    //  }

    //  vtkm::ErrorCode status = locator.FindCell(0.999 * p, cell_id, pcoords);
    //  if (cell_id >= 0)
    //  {
    //    auto cellShape = cell_set.GetCellShape(cell_id);
    //    auto indices = cell_set.GetIndices(cell_id);
    //    auto fieldValues = vtkm::make_VecFromPortalPermute(&indices, &field_in);
    //    vtkm::exec::CellInterpolate(fieldValues, pcoords, cellShape, out);
    //  }
    //  else
    //  {
    //    console::Error(point);
    //    this->RaiseError(" Given point outside of the cell set .");
    //  }
    //}
  }

  Vec3f offsets[5] = { { 0, 0, 0 },
                       { 0, 0.99999, 0 },
                       { 0, -0.99999, 0 },
                       { 0.99999, 0, 0 },
                       { -0.99999, 0, 0 } };
};

struct DispWorklet : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn coord_old,
                                FieldIn coord_new,
                                FieldOut pts_src,
                                FieldOut pts_tgt);
  using ExecutionSignature = void(_1, _2, _3, _4);

  DispWorklet(int current_step, int n_steps)
    : _current_step(current_step)
    , _n_steps(n_steps)
  {
  }

  template<typename T1, typename T2>
  VTKM_EXEC void operator()(const T1& coord_start,
                            const T2& coord_end,
                            Vec3f& pts_src,
                            Vec3f& pts_tgt) const
  {
    pts_src = coord_start + static_cast<Real>(_current_step) / _n_steps * (coord_end - coord_start);
    pts_tgt =
      coord_start + static_cast<Real>(_current_step + 1) / _n_steps * (coord_end - coord_start);
  }
  int _current_step;
  int _n_steps;
};

struct Remap : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature = void(CellSetIn cell_set_in,
                                FieldInPoint coord_src,
                                FieldInPoint coord_tgt,
                                FieldInCell cell_link,
                                WholeArrayIn value_src,
                                FieldOutCell value_tgt);

  using ExecutionSignature = void(InputIndex, _2, _3, _4, _5, _6);

  //using ScatterType = vtkm::worklet::ScatterUniform<3>;

  //Remap(int nstep)
  //  : _n_steps(0)
  //{
  //  dt = 1;
  //}

  VTKM_EXEC Vec3f Normal_Area(const Vec3f& ps, const Vec3f& pe) const
  {
    //Vec3f normal{ ps[1] - pe[1], pe[0] - ps[0], 0 };
    Vec3f normal{ pe[1] - ps[1], ps[0] - pe[0], 0 };
    //return vtkm::Normal(normal);
    return normal;
  }

  template<typename PType>
  VTKM_EXEC Real Area(const PType& pts) const
  {
    typename PType::ComponentType v1 = pts[1] - pts[0];
    typename PType::ComponentType v2 = pts[2] - pts[0];
    return 0.5 * (vtkm::Cross(v1, v2)[2]);
    //return (0.5) * (vtkm::Magnitude(vtkm::Cross(v1, v2)));
  }

  template<typename Pts, typename CellLink, typename Variable>
  VTKM_EXEC void operator()(Id& cell_id,
                            const Pts& pts_src,
                            const Pts& pts_tgt,
                            const CellLink& cell_link,
                            const Variable& value_src,
                            Real& value_tgt) const
  {
    Vec<Vec3f, 3> vel;
    for (IdComponent i = 0; i < 3; i++)
    {
      vel[i] = pts_tgt[i] - pts_src[i];
    }

    auto value = value_src.Get(cell_id);
    Real rhs = 0;
    Real rhs_area = 0;
    for (IdComponent face_index = 0; face_index < 3; face_index++)
    {
      IdComponent id0, id1;
      if (face_index == 0)
      {
        id0 = 0;
        id1 = 1;
      }
      if (face_index == 1)
      {
        id0 = 1;
        id1 = 2;
      }
      if (face_index == 2)
      {
        id0 = 2;
        id1 = 0;
      }

      Id neighbor_cell_id = cell_link[face_index].first;
      Real neighbor_value;
      if (neighbor_cell_id < 0)
      {
        //std::cout << cell_id << std::endl;
        neighbor_value = value;
      }
      else
      {
        neighbor_value = value_src.Get(neighbor_cell_id);
      }

      Vec3f normal = Normal_Area(pts_tgt[id0], pts_tgt[id1]);
      Vec3f face_vel = (vel[id0] + vel[id1]) / 2;
      Real vn = vtkm::Dot(face_vel, normal);

      //Real flux = 0.5 * (value + neighbor_value) * vn - 0.1*vtkm::Abs(vn) * (value - neighbor_value);
      Real flux = value * vn;
      if (vn > 0)
      {
        flux = neighbor_value * vn;
      }
      rhs += flux;
      rhs_area += vn;
    }

    Real area_src = Area(pts_src);
    Real area_tgt = Area(pts_tgt);
    Real area = (Area(pts_src) + dt * rhs_area);

    value_tgt = (value * Area(pts_src) + dt * rhs) / area;
    //if (cell_id == 0)
    //{
    //  std::cout << area_src << " " << area_tgt << " " << area << std::endl;
    //}
  }

  IdComponent _current_steps;
  Real dt = 1;
};

struct PeriodicMap : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature = void(CellSetIn cell_set_in,
                                FieldInPoint coord_src,
                                FieldInOutCell cell_link,
                                ExecObject locator);

  using ExecutionSignature = void(InputIndex, _2, _3, _4);

  template<typename Pts, typename CellLink, typename LocatorType>
  VTKM_EXEC void operator()(Id& cell_id,
                            const Pts& pts,
                            CellLink& cell_link,
                            const LocatorType& locator) const
  {
    for (IdComponent face_index = 0; face_index < 3; face_index++)
    {
      IdComponent id0, id1;
      if (face_index == 0)
      {
        id0 = 0;
        id1 = 1;
      }
      if (face_index == 1)
      {
        id0 = 1;
        id1 = 2;
      }
      if (face_index == 2)
      {
        id0 = 2;
        id1 = 0;
      }

      Id neighbor_cell_id = cell_link[face_index].first;
      if (neighbor_cell_id < 0)
      {
        Id find_id;
        Vec3f pcoords;
        IdComponent bid = -1;
        Vec3f p = (pts[id0] + pts[id1]) / 2.;
        Vec3f offset;
        if (vtkm::Abs(p[0] - _x_min) < _tol) // left
        {
          bid = 3;
          offset = Vec3f{ 1, 0, 0 };
        }
        else if (vtkm::Abs(p[0] - _x_max) < _tol) // right
        {
          bid = 1;
          offset = Vec3f{ -1, 0, 0 };
        }
        else if (vtkm::Abs(p[1] - _y_min) < _tol) // bottom
        {
          bid = 0;
          offset = Vec3f{ 0, 1, 0 };
        }
        else if (vtkm::Abs(p[1] - _y_max) < _tol) // top
        {
          bid = 2;
          offset = Vec3f{ 0, -1, 0 };
        }
        else // 内部点
        {
          bid = -1;
        }

        p += offset;
        vtkm::ErrorCode status = locator.FindCell(0.999 * p, find_id, pcoords);
        if (find_id >= 0)
        {
          //std::cout << cell_id << " " << find_id << std::endl;
          cell_link[face_index].first = find_id;
        }
        else
        {
          this->RaiseError(" Given point outside of the cell set .");
        }
      }
    }
  }
  vtkm::FloatDefault _x_min = 0, _x_max = 1, _y_min = 0, _y_max = 1;
  vtkm::FloatDefault _tol = 1E-06;
};

RemapUserObject::RemapUserObject(const Configuration& cfg)
  : UserObject(cfg)
  , _app(*(Get<Application*>("_app")))
  , _sys(_app.GetSystem<SGHSystem>())
  , _dtau(1)
  , _n_step(Get<int>("n_step"))
  , _L1(Get<Real>("L1", 1000))
  , _L2(Get<Real>("L2", 10))
  , _interval(Get<int>("interval", 1))
  , _periodic_bc(Get<bool>("periodic_bc", false))
{
}

void RemapUserObject::Init()
{
  console::Info("RemapUserObject::Init");

  ArrayHandle<Vec3f> pts;
  vtkm::cont::ArrayCopy(_app.GetMesh()->GetCoordinate().GetData(), pts);
  _coord = CoordinateSystem("coord_old", pts);

  vtkm::cont::CellLocatorTwoLevel locator;
  locator.SetDensityL1(_L1);
  locator.SetDensityL2(_L2);
  //vtkm::cont::CellLocatorGeneral locator;

  auto& cell_set = _sys.GetDateSet().GetCellSet();
  auto& coord = _sys.GetDateSet().GetCoordinateSystem();
  locator.SetCellSet(cell_set);
  locator.SetCoordinates(_coord);
  locator.Update();

  auto& mesh = *_app.GetMesh();
  auto& cell_links = mesh.GetFlatCellLinks();
  const auto& cell_links_offset = vtkm::cont::make_ArrayHandleGroupVec<3>(cell_links);
  Invoker{}(PeriodicMap{}, cell_set, _coord, cell_links_offset, locator);
}

void RemapUserObject::BuildPeriodicMap() {}

void RemapUserObject::Execute()
{
  auto current_step = static_cast<Transient&>(*(_app.GetExecutioner())).CurrentStep();


  if (current_step % _interval == 0)
  {
    RemapVarialbe();
    LocateVariable();
    _sys.InitVariable();
    _sys.PreSolve();
  }

  //ArrayHandle<Vec3f> disp;
  //vtkm::cont::Invoker{}(RemapWorklet{}, _coord, _sys.GetDateSet().GetCoordinateSystem(), disp);

  //auto energy = _sys.GetVariableArray<Real>("energy");
  //DataSet start, end;
  //start.AddCoordinateSystem(_coord);
  //start.SetCellSet(_sys.GetDateSet().GetCellSet());
  //start.AddCellField("energy", energy);
  //start.AddPointField("disp", disp);

  //vtkm::io::VTKDataSetWriter writer("start.vtk");
  //writer.WriteDataSet(start);
}

void RemapUserObject::RemapVarialbe()
{
  {
    Timer timer;
    timer.Start();
    auto& cell_set = _sys.GetDateSet().GetCellSet();
    auto& coord = _sys.GetDateSet().GetCoordinateSystem();
    auto& mesh = *_app.GetMesh();
    auto& cell_links = mesh.GetFlatCellLinks();

    auto& face_num_offset = mesh.GetFaceNumOffset();
    const auto& cell_links_offset = vtkm::cont::make_ArrayHandleGroupVec<3>(cell_links);
    //vtkm::cont::make_ArrayHandleGroupVecVariable(cell_links, face_num_offset);

    auto value_start = _sys.GetVariableArray<Real>("pressure");
    //ArrayHandle<Real> value_start;
    //value_start.Allocate(cell_set.GetNumberOfCells());
    //for (size_t i = 0; i < cell_set.GetNumberOfCells(); i++)
    //{
    //  value_start.WritePortal().Set(i, 1);
    //}
    ArrayHandle<Real> value_src;
    ArrayHandle<Real> value_tgt;
    vtkm::cont::ArrayCopy(value_start, value_src);
    for (int i = 0; i < _n_step; i++)
    {
      console::Info("step: ", i);
      ArrayHandle<Vec3f> pts_src, pts_tgt;
      Invoker{}(DispWorklet{ i, _n_step }, coord, _coord, pts_src, pts_tgt);
      Invoker{}(Remap{}, cell_set, pts_src, pts_tgt, cell_links_offset, value_src, value_tgt);
      vtkm::cont::ArrayCopy(value_tgt, value_src);
      //value_src = value_tgt;
    }
    console::Info("RemapUserObject::Remap pressure ", timer.GetElapsedTime());
    vtkm::cont::ArrayCopy(value_tgt, value_start);
  }

  {
    Timer timer;
    timer.Start();
    auto& cell_set = _sys.GetDateSet().GetCellSet();
    auto& coord = _sys.GetDateSet().GetCoordinateSystem();
    auto& mesh = *_app.GetMesh();
    auto& cell_links = mesh.GetFlatCellLinks();

    auto& face_num_offset = mesh.GetFaceNumOffset();
    const auto& cell_links_offset = vtkm::cont::make_ArrayHandleGroupVec<3>(cell_links);
    //vtkm::cont::make_ArrayHandleGroupVecVariable(cell_links, face_num_offset);

    auto value_start = _sys.GetVariableArray<Real>("density");
    //ArrayHandle<Real> value_start;
    //value_start.Allocate(cell_set.GetNumberOfCells());
    //for (size_t i = 0; i < cell_set.GetNumberOfCells(); i++)
    //{
    //  value_start.WritePortal().Set(i, 1);
    //}
    ArrayHandle<Real> value_src;
    ArrayHandle<Real> value_tgt;
    vtkm::cont::ArrayCopy(value_start, value_src);
    for (int i = 0; i < _n_step; i++)
    {
      console::Info("step: ", i);
      ArrayHandle<Vec3f> pts_src, pts_tgt;
      Invoker{}(DispWorklet{ i, _n_step }, coord, _coord, pts_src, pts_tgt);
      Invoker{}(Remap{}, cell_set, pts_src, pts_tgt, cell_links_offset, value_src, value_tgt);
      vtkm::cont::ArrayCopy(value_tgt, value_src);
      //value_src = value_tgt;
    }
    console::Info("RemapUserObject::Remap density ", timer.GetElapsedTime());
    vtkm::cont::ArrayCopy(value_tgt, value_start);
  }
  //DataSet target;
  //target.AddCoordinateSystem(_coord);
  //target.SetCellSet(_sys.GetDateSet().GetCellSet());
  //target.AddCellField("target", value_tgt);

  //vtkm::io::VTKDataSetWriter writer_tgt("target.vtk");
  //writer_tgt.WriteDataSet(target);

  //DataSet source;
  //source.AddCoordinateSystem(_coord);
  //source.SetCellSet(_sys.GetDateSet().GetCellSet());
  //source.AddCellField("source", value_src);
  //vtkm::io::VTKDataSetWriter writer_src("source.vtk");
  //writer_src.WriteDataSet(source);
}
void RemapUserObject::LocateVariable()
{
  Timer timer;
  timer.Start();
  ArrayHandle<Id> cell_ids;
  ArrayHandle<Vec3f> pts;
  //vtkm::cont::CellLocatorBoundingIntervalHierarchy locator;
  //locator.SetNumberOfSplittingPlanes(10);
  //locator.SetMaxLeafSize(10);
  vtkm::cont::CellLocatorTwoLevel locator;
  locator.SetDensityL1(_L1);
  locator.SetDensityL2(_L2);
  //vtkm::cont::CellLocatorGeneral locator;

  auto& cell_set = _sys.GetDateSet().GetCellSet();
  auto& coord = _sys.GetDateSet().GetCoordinateSystem();
  locator.SetCellSet(cell_set);
  locator.SetCoordinates(coord);
  locator.Update();

  const auto& velocity = _sys.GetVariableArray<Vec3f>("velocity");
  ArrayHandle<Vec3f> vel_out;

  if (_periodic_bc)
  {
    Invoker{}(InterpolateVelocity{}, _coord, locator, cell_set, velocity, vel_out);
  }
  else
  {
    Invoker{}(FindCellWorklet{}, _coord, locator, cell_set, velocity, vel_out);
  }

  vtkm::cont::ArrayCopy(vel_out, velocity);
  vtkm::cont::ArrayCopy(_coord.GetData(), coord.GetData());
  //console::Info("RemapUserObject::LocateVariable", timer.GetElapsedTime());

  //DataSet source;
  //source.AddCoordinateSystem(_coord);
  //source.SetCellSet(_sys.GetDateSet().GetCellSet());
  //source.AddPointField("vel", vel_out);
  //vtkm::io::VTKDataSetWriter writer_src("vel.vtk");
  //writer_src.WriteDataSet(source);
}
